# -*- coding: utf-8 -*-

"""
DateTime   : 2021/04/06 11:42
Author     : ZhangYafei
Description: 
"""
import os
from concurrent.futures import ProcessPoolExecutor

import requests
from prettytable import PrettyTable
from tqdm import tqdm

from utils import print_color

image_type_info = {
    0: '最新壁纸', 6: '美女模特', 30: '爱情美图', 9: '风景大片',
    36: '4K专区', 15: '小清新', 26: '动漫卡通', 11: '明星风尚',
    14: '萌宠动物', 5: '游戏壁纸', 12: '汽车天下', 10: '炫酷时尚',
    29: '月历壁纸', 7: '景视剧照', 35: '文字控', 13: '节日美图',
    22: '军事天地', 16: '劲爆体育', 18: 'BABY秀'
}

image_data_list = set()
history_file_path = 'history.txt'
history_file = open(history_file_path, mode='a+', encoding='utf-8')
downloaded_images = set()
if os.path.exists(history_file_path):
    history_file.seek(0)
    downloaded_images = {line.strip() for line in history_file}


def init_tips_info():
    table = PrettyTable(['序号1', '壁纸分类1', '序号2', '壁纸分类2'])
    row = []
    for index, key in enumerate(image_type_info, start=1):
        row.extend([key, image_type_info[key]])
        if len(row) == 4:
            table.add_row(row)
            row.clear()
            continue
        if index == len(image_type_info):
            row.extend(['-', '-'])
            table.add_row(row)
    print(table)


def search_images(image_type: int, start: int = 1, count: int = 50):
    print_color(f'正在检索 - {image_type_info[image_type]}')
    url = 'https://bird.ioliu.cn/v2'
    if image_type == 0:
        params = {
            'url': f'http://wallpaper.apc.360.cn/index.php?c=WallPaper&start={start}&count={count}&from=360chrome&a=getAppsByOrder&order=create_time'}
    else:
        params = {
            'url': f'http://wallpaper.apc.360.cn/index.php?c=WallPaper&start={start}&count={count}&from=360chrome&a=getAppsByCategory&cid={image_type}'}

    # 图片文件夹分类
    image_type_desc = image_type_info[image_type]
    if not os.path.exists(f'images/{image_type_desc}'):
        os.makedirs(f'images/{image_type_desc}')

    # 发送请求
    response = requests.get(url=url, params=params)
    result = response.json()
    downloaded_count = 0
    for image_info in result['data']:
        uid = image_info['id']
        image_url = image_info['url']
        tags = image_info['tag'].replace('_category_', '').replace('_全部_', '').strip('_ ').replace(' ', '')
        size = image_info['resolution'].replace('x', '_')
        filepath = f'images/{image_type_info[image_type]}/{tags}_{uid}_{size}.jpg'
        if filepath not in downloaded_images:
            image_data_list.add((image_url, filepath))
        else:
            downloaded_count += 1
    return image_data_list, downloaded_count


def start_download(image_data_list):
    if len(image_data_list) > 30:
        choice = input('是否开启多进程加速下载？(Y/N) >> ')
        if choice.upper() == 'Y':
            data_list = [set() for _ in range(4)]
            for index, image in enumerate(image_data_list):
                remainder = index % 4
                data_list[remainder].add(image)
            with ProcessPoolExecutor(max_workers=4) as pool:
                pool.map(download_images, data_list)
            print_color(f'\n\n所有图片下载完成!')
            return
    download_images()
    print_color(f'所有图片下载完成!')


def download_images(data_list: set = None):
    if data_list:
        task_progress = tqdm(data_list, ncols=100)
    else:
        task_progress = tqdm(image_data_list, ncols=100)
    for index, task in enumerate(task_progress, start=1):
        url, filepath = task
        task_progress.set_description(f"正在下载第 {index} 张图片")
        response = requests.get(url=url)
        with open(filepath, mode='wb') as f:
            f.write(response.content)
        history_file.write(f'{filepath}\n')


def run():
    init_tips_info()
    while True:
        image_type = input('请选择下载的壁纸分类(序号,默认回车为0) >>  ')
        if image_type:
            if not image_type.isdecimal() or int(image_type) not in image_type_info:
                print('您输入的分类序号有误，请重新确认后输入')
                continue
            else:
                image_type = int(image_type)
        else:
            image_type = 0

        break
    while True:
        count = input('请输入下载数量（整数，默认为20）>> ')
        if count:
            if not count.isdecimal():
                print('您输入的数量有误，必须为整数，请重新输入')
                continue
            else:
                count = int(count)
        else:
            count = 20
        break

    image_data_list, downloaded_count = search_images(image_type=image_type, count=count)

    if downloaded_count == 0:
        print_color(f'共检索到 {len(image_data_list)} 张图片, 开始下载')
    elif len(image_data_list) == 0:
        print_color(f'共检索到 {downloaded_count} 张图片, 已全部下载完成！')
        return
    else:
        print_color(
            f'共检索到 {len(image_data_list) + downloaded_count} 张图片, 已下载 {downloaded_count} 张, 还需下载 {len(image_data_list)} 张, 开始下载')

    start_download(image_data_list=image_data_list)


if __name__ == '__main__':
    run()
